"""
    Data simulation script
    @author: Shao-Chuan Wang
    @Date: Apr, 06, 2011
"""
import random
import sys
import math

OUTPUT_FILENAME = 'foo.sto'
INIT_PRICE = 10  # initial price
NTICKS = 100000     # number of total ticks


def go_up(curr_price, ratio):
  db = curr_price / INIT_PRICE
  ratio -= 0.005 * math.log(db)
  inst = random.randint(0,100)
  return inst < int(100 * ratio)

def simu_prices(INIT_PRICE, prob_up_range, delta_range, NTICKS):  
  prices = [(0, INIT_PRICE, INIT_PRICE, INIT_PRICE, INIT_PRICE)]
  volumns = [10]

  for t in xrange(NTICKS):
    if (t % 10 == 0):
      sys.stdout.write('.')
      sys.stdout.flush()
    prob_up = prob_up_range[0] + random.random() * (prob_up_range[1]-prob_up_range[0])
    delta = random.random() * delta_range
    candidates = [random.random() * 2 * delta_range - delta_range,
        random.random() * 2 * delta_range - delta_range,
        random.random() * 2 * delta_range - delta_range,
        random.random() * 2 * delta_range - delta_range,
        delta,
        -delta
    ]
    open_ = candidates[0]
    high = max(candidates)
    low = min(candidates)
    volumns.append( 1 + int(random.random() * 100) + \
            int(10000.0 * abs(high) * random.random() * random.random()))
    last_close = prices[-1][2]
    if go_up(last_close, prob_up):
      p = (t+1, 
              last_close*(1+open_), 
              last_close*(1+delta), 
              last_close*(1+high), 
              last_close*(1+low))
      prices.append(p)
    else:
      p = (t+1, 
              last_close*(1+open_),
              last_close*(1-max(delta-0.0024,0)), 
              last_close*(1+high),
              last_close*(1+low))
      prices.append(p)

  return prices,volumns

def dump_stock(prices, volumns, stockname):
    fp = open(stockname, 'w')
    print >>fp, '%%t\topen\tclose\thigh\tlow\tvolumn'
    for pr, vol in zip(prices, volumns):
        t,o,c,h,l = pr
        print >> fp, '%d\t%f\t%f\t%f\t%f\t%d' % (t,o,c,h,l,vol)
    fp.close()

def plot_candlestick(prices):
    import matplotlib.pyplot as plt
    import matplotlib.finance as finance
    # Plotting
    fig = plt.figure(facecolor='white')
    ax = fig.add_subplot(111)

def main():
  prob_up_range = (0.4, 0.6)  # winning priors
  delta_range = 0.05  # up or down 10% upmost

  prices,volumns= simu_prices(INIT_PRICE, prob_up_range, delta_range, NTICKS)
  dump_stock(prices, volumns, OUTPUT_FILENAME)
  print "max: ", max([p[1] for p in prices])
  return

if __name__=='__main__':
  main()
  raw_input()
